#coding:utf-8
from ultralytics import YOLO
import cv2

# 所需加载的模型目录
#path = 'models/best.pt'
# 需要检测的图片地址
#img_path = "TestFiles/IM-0003-0001.jpeg"

# 加载模型
#model = YOLO(path, task='classify')

# 检测图片
#results = model(img_path)
#print(results)
#res = results[0].plot()
#res = cv2.resize(res,dsize=None,fx=0.3,fy=0.3,interpolation=cv2.INTER_LINEAR)
#cv2.imshow("YOLOv8 Detection", res)
#cv2.waitKey(0)
import cv2
from ultralytics import YOLO
import os
from datetime import datetime

# 模型路径
model_path = 'models/best.pt'
# 待分类图片路径
img_path = "TestFiles/IM-0003-0001.jpeg"

# 创建保存结果的目录（如果不存在）
output_dir = "results"
os.makedirs(output_dir, exist_ok=True)

# 加载分类模型
model = YOLO(model_path, task='classify')

# 运行分类
results = model(img_path)

# 获取分类概率
probs = results[0].probs
top5_classes = probs.top5
top5_conf = probs.top5conf

# 打印结果
print("Top 5 classes:", top5_classes)
print("Confidences:", top5_conf)

# 可视化结果（带标签的图片）
res_img = results[0].plot()

# 调整图片大小（可选）
res_img = cv2.resize(res_img, None, fx=0.3, fy=0.3)

# 生成唯一文件名（防止覆盖）
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_img_path = os.path.join(output_dir, f"result_{timestamp}.jpg")
output_txt_path = os.path.join(output_dir, f"result_{timestamp}.txt")

# 保存图片
cv2.imwrite(output_img_path, res_img)

# 保存分类结果到TXT文件（可选）
with open(output_txt_path, 'w') as f:
    f.write(f"Image: {img_path}\n")
    f.write("Classification Results:\n")
    for i, (class_id, conf) in enumerate(zip(top5_classes, top5_conf)):
        f.write(f"{i+1}. Class: {class_id}, Confidence: {conf:.4f}\n")

# 显示结果（可选）
cv2.imshow("YOLOv8 Classification", res_img)
cv2.waitKey(0)
cv2.destroyAllWindows()

print(f"Results saved to: {output_img_path} and {output_txt_path}")